import numpy as np
import torch
import torch.backends.cudnn as cudnn

from code.net.facenet.Facenet import Facenet as facenet
from code.utils.utils import preprocess_input, resize_image


# --------------------------------------------#
#   使用自己训练好的模型预测需要修改2个参数
#   model_path和backbone需要修改！
# --------------------------------------------#
class Facenet(object):

    # ---------------------------------------------------#
    #   初始化Facenet
    # ---------------------------------------------------#
    def __init__(self, model_path, **kwargs):
        self.model_path = model_path
        # 添加默认网络参数
        defaults = {
            #   输入图片的大小。
            # --------------------------------------------------------------------------#
            "input_shape": [160, 160, 3],
            # --------------------------------------------------------------------------#
            #   所使用到的主干特征提取网络
            # --------------------------------------------------------------------------#
            "backbone": "mobilenet",
            # -------------------------------------------#
            #   是否进行不失真的resize
            # -------------------------------------------#
            "letterbox_image": True,
            # -------------------------------------------#
            #   是否使用Cuda
            #   没有GPU可以设置成False
            # -------------------------------------------#
            "cuda": False,
            "threshold": 0.6
        }
        self.__dict__.update(defaults)
        for name, value in kwargs.items():
            setattr(self, name, value)

        self.generate()       # 生成网络


    # 网络的生成
    def generate(self):
        # ---------------------------------------------------#
        #   载入模型与权值
        # ---------------------------------------------------#

        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')        # 这里采用cpu
        self.net = facenet(backbone=self.backbone).eval()            # 这里使用eval()而不是train 因为是预测 采用的是训练好的
        self.net.load_state_dict(torch.load(self.model_path, map_location=device), strict=False)      # 这些self的属性都在__init__中生成 这里不是报错

        if self.cuda:
            self.net = torch.nn.DataParallel(self.net)
            cudnn.benchmark = True
            self.net = self.net.cuda()

    # ---------------------------------------------------#
    #   检测图片
    # ---------------------------------------------------#
    def detect_image(self, image_1, image_2):
        # ---------------------------------------------------#
        #   图片预处理，归一化
        # ---------------------------------------------------#
        with torch.no_grad():
            image_1 = resize_image(image_1, [self.input_shape[1], self.input_shape[0]],
                                   letterbox_image=self.letterbox_image)
            image_2 = resize_image(image_2, [self.input_shape[1], self.input_shape[0]],
                                   letterbox_image=self.letterbox_image)

            photo_1 = torch.from_numpy(
                np.expand_dims(np.transpose(preprocess_input(np.array(image_1, np.float32)), (2, 0, 1)), 0))       # transpose将图片numpy数组转变为通道在前  expand_dims将batch size添加上(1)
            photo_2 = torch.from_numpy(
                np.expand_dims(np.transpose(preprocess_input(np.array(image_2, np.float32)), (2, 0, 1)), 0))

            if self.cuda:
                photo_1 = photo_1.cuda()
                photo_2 = photo_2.cuda()

            # ---------------------------------------------------#
            #   图片传入网络进行预测
            # ---------------------------------------------------#
            output1 = self.net(photo_1).cpu().numpy()
            output2 = self.net(photo_2).cpu().numpy()

            # ---------------------------------------------------#
            #   计算二者之间的距离
            # ---------------------------------------------------#
            distance = np.linalg.norm(output1 - output2, axis=1)

        return distance

    # 获取相似概率
    def similar_probability(self, image_1, image_2):
        distance = self.detect_image(image_1, image_2)     # 获得距离
        # probability = 0.5 - (distance - self.threshold)/self.threshold
        return distance
        # if probability < 0:
        #     return 0
        # elif probability > 1:
        #     return 1
        # else:
        #     return np.round(probability,4)





